import datetime
import pymysql
import pandas as pd
from sklearn.linear_model import LinearRegression
import warnings

warnings.filterwarnings("ignore")


def read_sql(sql, db_name):
    db_config = {
        "charset": "utf8",
        "host": "192.168.0.114",
        "port": 3307,
        "user": "admin",
        "passwd": "123456"
    }
    conn = pymysql.connect(db=db_name, **db_config)
    return pd.read_sql(sql, conn)


def fix_fill_times(binance_df, fill_times, expand_kline_list, step):
    """
    binance缺失数据之前的一条数据有可能会出现vol=0,o=h=l=c， 这样的数据也是无效数据需要替换
    :param binance_df:
    :param fill_times:
    :param expand_kline_list:
    :param step:
    :return:
    """
    for i, fill_time in enumerate(fill_times):
        before_missing_time = fill_time[0] - datetime.timedelta(seconds=step)
        if binance_df[binance_df["datetime"] == before_missing_time].iloc[0]["vol"] == 0:
            print(
                f"{before_missing_time}存在无效数据:{binance_df[binance_df["datetime"] == before_missing_time].iloc[0].to_dict()}")
            fill_times[i].insert(0, before_missing_time)
            expand_kline_list.append(before_missing_time)
    return fill_times


def get_fill_times(binance_df, else_df, frequency):
    """
    得到需要填充的时间段
    :param binance_df:
    :param else_df: 其他数据源数据
    :param frequency: 1h或5m
    :return: 需要用其他数据源填充的时间段以及binance存在但是无效的时间
    """
    #  把连续的缺失分为1组
    miss_times = list(set(else_df["datetime"]) - set(binance_df["datetime"]))
    miss_times.sort()
    fill_times = [[miss_times[0]]]
    expand_kline_list = []  # binance存在但是无效的数据
    index = 0
    if frequency == "1h":
        step = 3600
    elif frequency == "5m":
        step = 300
    else:
        raise Exception("invalid frequency")
    for i in range(1, len(miss_times)):
        if (miss_times[i] - miss_times[i - 1]).total_seconds() == step:
            fill_times[index].append(miss_times[i])
        else:
            expand_one_step = fill_times[index][-1] + datetime.timedelta(seconds=step)
            fill_times[index].append(expand_one_step)
            expand_kline_list.append(expand_one_step)
            index += 1
            fill_times.append([])
            fill_times[index].append(miss_times[i])
    expand_one_step = fill_times[index][-1] + datetime.timedelta(seconds=step)
    expand_kline_list.append(expand_one_step)
    fill_times[index].append(expand_one_step)
    fix_fill_times(binance_df, fill_times, expand_kline_list, step)
    return fill_times, expand_kline_list


def fill_binance_missing_with_else_datasource(binance_df, else_df, frequency, datasource):
    assert frequency != "1d", "暂不支持填充缺失日线"
    cond1 = else_df["datetime"] <= binance_df["datetime"].max()
    cond2 = else_df["datetime"] >= binance_df["datetime"].min()
    else_df = else_df[cond1 & cond2]
    fill_times, expand_kline_list = get_fill_times(binance_df, else_df, frequency)
    tmp_else_df_list = []
    print(f"从{datasource}填充{len(fill_times)}段数据")
    count = 0
    for fill_time_list in fill_times:
        count += 1
        # print(f"第{count}段,填充开始时间:{fill_time_list[0]},填充结束时间:{fill_time_list[-1]}")
        tmp_else_df = else_df[else_df["datetime"].isin(fill_time_list)]
        tmp_else_df_list.append(tmp_else_df)
    fill_else_df = pd.concat(tmp_else_df_list, ignore_index=True)
    binance_df = binance_df[~binance_df["datetime"].isin(expand_kline_list)]
    binance_df = pd.concat([binance_df, fill_else_df], ignore_index=True)
    binance_df.sort_values(by=["datetime"], inplace=True)
    binance_df.reset_index(drop=True, inplace=True)
    print(f"从{datasource}填充的数据长度:{len(fill_else_df)}")
    if datasource == "bitstamp" or datasource == "glassnode":  # bitstamp或glassnode的数据都是完备的 可以用来辅助检查一下binance的数据
        else_df.sort_values(by=["datetime"], inplace=True)
        assert binance_df["datetime"].reset_index(drop=True).equals(else_df["datetime"].reset_index(drop=True))
    return binance_df


def fill_missing_vol(binance_df):

    null_df = binance_df[binance_df["vol"].isna()]
    null_df = null_df.reset_index(drop=False)
    null_df['group'] = (null_df['index'].diff() != 1).cumsum()
    fill_values = []
    for grp, grp_null_df in null_df.groupby("group"):
        first_time = grp_null_df.iloc[0]["datetime"]
        last_time = grp_null_df.iloc[-1]["datetime"]
        if sum(pd.notna(binance_df[binance_df["datetime"] < first_time]["okx_vol"])) < len(grp_null_df)*5:
            print(f"{first_time}前okx没有数据, 需要用bitsatmp填充")
            df_temp = linear_regression_fill(binance_df[binance_df["datetime"] <= last_time], "bitstamp_vol", "vol", len(grp_null_df)*6)
            fill_values.extend(df_temp[-len(grp_null_df):].tolist())
        else:
            print(f"{first_time}前okx有数据, 用okx填充")
            df_temp = linear_regression_fill(
                binance_df[(binance_df["datetime"] <= last_time) & (binance_df['okx_vol'].notna())], "okx_vol", "vol", len(grp_null_df)*6)
            fill_values.extend(df_temp[-len(grp_null_df):].tolist())
    binance_df.loc[binance_df["vol"].isna(), "vol"] = fill_values
    return binance_df


def get_binance_kline_with_vol(frequency="1h"):
    """
    得到币安的k线(包含成交量)，开始时间为币安有数据开始
    缺失的价格部分用okx或glassnode或bitstamp填充
    缺失的成交量部分暂时空着
    :param frequency:
    :return:
    """
    binance_df = read_sql(f"select * from k_line where symbol='BTC' and frequency='{frequency}' "
                          f"and datasource = 'binance' order by datetime", db_name="all_history_ohlcvm_coinmarketcap")
    if frequency in ["1h", "5m"]:
        binance_begin_time = binance_df["datetime"].iloc[0]
        okx_df = read_sql(
            f"select * from k_line where symbol='BTC' and frequency='{frequency}' and datasource = 'okx' "
            f"and datetime >= '{binance_begin_time}' order by datetime", db_name="all_history_ohlcvm_coinmarketcap")
        okx_df_without_vol = okx_df.drop(columns=['vol', 'amount'])  # 填充缺失值只填充价格
        binance_df = fill_binance_missing_with_else_datasource(binance_df, okx_df_without_vol, frequency,
                                                               datasource="okx")
        if frequency == "1h":
            glassnode_df = read_sql(
                f"select * from k_line where symbol='BTC' and frequency='{frequency}' "
                f"and datasource = 'glassnode' and datetime >= '{binance_begin_time}' order by datetime",
                db_name="all_history_ohlcvm_coinmarketcap")
            glassnode_df_without_vol = glassnode_df.drop(columns=['vol', 'amount'])  # 填充缺失值只填充价格
            binance_df = fill_binance_missing_with_else_datasource(binance_df, glassnode_df_without_vol, frequency,
                                                                   datasource="glassnode")
            bitstamp_df = read_sql(
                f"select * from k_line where symbol='BTC' and frequency='{frequency}' "
                f"and datasource = 'bitstamp' and datetime >= '{binance_begin_time}' order by datetime",
                db_name="all_history_ohlcvm_coinmarketcap")
        else:
            bitstamp_df = read_sql(
                f"select * from k_line where symbol='BTC' and frequency='{frequency}' "
                f"and datasource = 'bitstamp' and datetime >= '{binance_begin_time}' order by datetime",
                db_name="all_history_ohlcvm_coinmarketcap")
            bitstamp_df_without_vol = bitstamp_df.drop(columns=['vol', 'amount'])  # 填充缺失值只填充价格
            binance_df = fill_binance_missing_with_else_datasource(binance_df, bitstamp_df_without_vol, frequency,
                                                                   datasource="bitstamp")
        binance_df = pd.merge(binance_df, okx_df[["datetime", "vol", "amount"]].rename(columns={"vol": "okx_vol",
                                                                                                "amount": "okx_amount"}),
                              on="datetime", how="left")
        binance_df = pd.merge(binance_df, bitstamp_df[["datetime", "vol"]].rename(columns={"vol": "bitstamp_vol"}),
                              on="datetime", how="left")
        columns = ["datetime", "open", "high", "low", "close", "vol", "amount", "okx_vol", "okx_amount", "bitstamp_vol"]
        binance_df = binance_df[columns]
        binance_df.sort_values(by=["datetime"], inplace=True)
        binance_df = fill_missing_vol(binance_df)
    return binance_df[["datetime", "open", "high", "low", "close", "vol"]]


def linear_regression_fill(df, column_a, column_b, truncate_length):
    """
    使用线性回归填充缺失值
    df: 数据框
    column_a: 自变量列名
    column_b: 因变量列名（包含缺失值）
    """
    # 复制数据避免修改原数据
    df_temp = df[[column_a, column_b]].copy()
    df_temp = df_temp[-truncate_length:]  # 取缺失数据6倍长度
    # 分离有缺失值和无缺失值的数据
    known_data = df_temp[df_temp[column_b].notna()]
    unknown_data = df_temp[df_temp[column_b].isna()]

    if len(known_data) < 2:
        print("已知数据不足，无法建立回归模型")
        return df[column_b]

    # 准备训练数据
    X_train = known_data[[column_a]].values
    y_train = known_data[column_b].values

    # 训练线性回归模型
    model = LinearRegression()
    model.fit(X_train, y_train)

    # 预测缺失值
    if len(unknown_data) > 0:
        X_pred = unknown_data[[column_a]].values
        y_pred = model.predict(X_pred)

        # 填充缺失值
        df_temp.loc[df_temp[column_b].isna(), column_b] = y_pred

        print(f"使用线性回归填充了 {len(unknown_data)} 个缺失值")
        print(f"回归方程: y = {model.coef_[0]:.4f} * x + {model.intercept_:.4f}")
        print(f"R² 分数: {model.score(X_train, y_train):.4f}")

    return df_temp[column_b]


def get_null_with_context(df, column, before=1, after=1):
    """
    获取指定列为空值的行及其前后几行

    Parameters:
    df: DataFrame
    column: 要检查的列名
    before: 空值前的行数
    after: 空值后的行数
    """
    # 找到空值的索引位置
    null_indices = df[df[column].isna()].index

    # 收集所有需要包含的索引
    all_indices = set()

    for idx in null_indices:
        # 添加空值行本身
        all_indices.add(idx)

        # 添加前面的行
        for i in range(1, before + 1):
            prev_idx = idx - i
            if prev_idx >= df.index.min():
                all_indices.add(prev_idx)

        # 添加后面的行
        for i in range(1, after + 1):
            next_idx = idx + i
            if next_idx <= df.index.max():
                all_indices.add(next_idx)

    # 按原始顺序排序并返回
    result_indices = sorted(all_indices)
    return df.loc[result_indices]


def fill_missing_vol_1h():
    binance_df = pd.read_csv("BTC_kline_1h_with_vol.csv")
    binance_df["datetime"] = pd.to_datetime(binance_df["datetime"])
    tmp_binance_df = binance_df[["datetime", "open", "high", "low", "close", "vol", "amount"]]
    null_df = tmp_binance_df[tmp_binance_df.isna().any(axis=1)]
    null_df = null_df.reset_index(drop=False)
    null_df['group'] = (null_df['index'].diff() != 1).cumsum()
    binance_df["vol_ma700"] = binance_df["vol"].rolling(700, min_periods=100).mean()
    binance_df["amount_ma700"] = binance_df["amount"].rolling(700, min_periods=100).mean()
    fill_values = []
    for grp, grp_null_df in null_df.groupby("group"):
        first_time = grp_null_df.iloc[0]["datetime"]
        last_time = grp_null_df.iloc[-1]["datetime"]
        if sum(pd.notna(binance_df[binance_df["datetime"] < first_time]["okx_vol"])) < 100:
            print(f"{first_time}前okx没有数据, 需要用bitsatmp填充")
            df_temp = linear_regression_fill(binance_df[binance_df["datetime"] <= last_time], "bitstamp_vol", "vol", len(grp_null_df)*6)
            fill_values.extend(df_temp[-len(grp_null_df):].tolist())
        else:
            print(f"{first_time}前okx有数据, 用okx填充")
            df_temp = linear_regression_fill(
                binance_df[(binance_df["datetime"] <= last_time) & (binance_df['okx_vol'].notna())], "okx_vol", "vol", len(grp_null_df)*6)
            fill_values.extend(df_temp[-len(grp_null_df):].tolist())
    binance_df.loc[binance_df["vol"].isna(), "vol_by_regression"] = fill_values
    binance_df.to_csv("BTC_kline_1h_with_vol_fill_missing.csv", index=False)
    binance_df_filter = get_null_with_context(binance_df, column="vol", before=10, after=10)
    binance_df_filter.to_csv("BTC_kline_1h_with_vol_fill_missing(截取拼接处).csv", index=False)


if __name__ == '__main__':
    BTC_kline_1d = get_binance_kline_with_vol(frequency="1d")
    BTC_kline_1d.to_csv("BTC_kline_1d_with_vol.csv", index=False)
    BTC_kline_1h = get_binance_kline_with_vol(frequency="1h")
    BTC_kline_1h.to_csv("BTC_kline_1h_with_vol.csv", index=False)
    BTC_kline_5m = get_binance_kline_with_vol(frequency="5m")
    BTC_kline_5m.to_csv("BTC_kline_5m_with_vol.csv", index=False)


